Self-balancing BST [AVL tree]

AVL tree (Adelson-Velsky and Landis) is a self-balancing binary search tree. In an AVL tree, the heights of the two child subtrees of any node differ by at most one; if at any time they differ by more than one, rebalancing is done to restore this property.

Lookup, insertion, and deletion all take O(logN) time in both the average and worst cases, where N is the number of nodes in the tree prior to the operation. Insertions and deletions may require the tree to be rebalanced by one or more tree rotations. AVL trees are height-balanced.

class AVLNode(object):

    def __init__(self, data, parentNode):
        self.data = data
        self.parentNode = parentNode
        self.left = None
        self.right = None
        self.balance = 0

    def insertAVL(self, data, parentNode):
        if data < self.data:
            if not self.left:
                self.left = AVLNode(data, parentNode)
            else:
                self.left.insertAVL(data, parentNode)
        else:
            if not self.right:
                self.right = AVLNode(data, parentNode)
            else:
                self.right.insertAVL(data, parentNode)
        # to check 1<balance<-1 -> make rotation
        return parentNode


    def traverseInOrderAVL(self):
        if self.left is not None:
            self.left.traverseInOrderAVL()
        print(self.data)
        if self.right is not None:
            self.right.traverseInOrderAVL()


    def getMinAVL(self):
        if self.left is None:
            return self.data
        else:
            return self.left.getMinAVL()

    def getMaxAVL(self):
        if self.right is None:
            return self.data
        else:
            return self.right.getMaxAVL()

class BalancedTree(object):

    def __init__(self):
        self.root = None

    def insert(self, data):

        parentNode = self.root

        if self.root == None:
            parentNode = AVLNode(data, None)
            self.root = parentNode
        else:
            parentNode = self.root.insertAVL(data, self.root)

        self.rebalanceTree(parentNode)

    def height(self, node):
        if node is None:
            return -1
        else:
            return max(self.height(node.left), self.height(node.right))

    def rebalanceTree(self, parentNode):

        self.setBalance(parentNode)

        # 4 situations with rotations
        if parentNode.balance <= -1:
            if self.height(parentNode.left.left) >= self.height(parentNode.left.right):
                parentNode = self.rotateRight(parentNode)
            else:
                parentNode = self.rotateLeftRight(parentNode)
        elif parentNode.balance > 1:
            if self.height(parentNode.right.right) >= self.height(parentNode.right.left):
                parentNode = self.rotateLeft(parentNode)
            else:
                parentNode = self.rotateRightLeft(parentNode)

        if parentNode.parentNode is not None:
            self.rebalanceTree(parentNode.parentNode)
        else:
            self.root = parentNode

    def rotateLeftRight(self, node):
        print("Rotation Left-Right...")
        node.left = self.rotateLeft(node.left)
        return self.rotateRight(node)

    def rotateRightLeft(self, node):
        print("Rotation Right-Left...")
        node.right = self.rotateRight(node.right)
        return self.rotateLeft(node)

    def setBalance(self, node):
        node.balance = self.height(node.right) - self.height(node.left)

    def rotateLeft(self, node):
        print("Rotation Left...")

        tmp = node.right
        tmp.parentNode = node.parentNode
        node.right = tmp.left

        if node.right is not None:
            node.right.parentNode = node

        tmp.left = node
        node.parentNode = tmp

        if tmp.parent is not None:
            if tmp.parentNode.right == node:
                tmp.parentNode.right = tmp
            else:
                tmp.parentNode.left = tmp

        self.setBalance(node)
        self.setBalance(tmp)   # after rotation it could be unbalanced

        return tmp

    def rotateRight(self, node):
        print("Rotation Right...")

        tmp = node.left
        tmp.parentNode = node.parentNode
        node.left = tmp.right

        if node.left is not None:
            node.left.parentNode = node

        tmp.right = node
        node.parentNode = tmp

        if tmp.parentNode is not None:
            if tmp.parentNode.left == node:
                tmp.parentNode.left = tmp
            else:
                tmp.parentNode.right = tmp

        self.setBalance(node)
        self.setBalance(tmp)   # after rotation it could be unbalanced

        return tmp

tree = BalancedTree()
# righ-left rotation
tree.insert(4)    # root
tree.insert(6)    # right   (6 > 4)
tree.insert(5)    # right.left (5 < 6)

root = tree.root
assert root.data == 4
assert root.right.data == 6
assert root.right.left.data == 5

rr = tree.rotateRight(root.right)

assert rr.data == 5
assert rr.right.data == 6
assert rr.parentNode.data == 4

Rotation Right...